"""
python fujian1_cubicSpline_fix.py
"""

import os
import json
import numpy as np
import pandas as pd
from scipy.interpolate import CubicSpline
import matplotlib.pyplot as plt

# 输入和输出路径
input_dir = "fujian/fujian1/group_output"
output_dir = "fujian/fujian1/cubic_spline/interpolation_output"
log_dir = "fujian/fujian1/cubic_spline/log"
graph_dir = "fujian/fujian1/cubic_spline/graph"
os.makedirs(output_dir, exist_ok=True)
os.makedirs(log_dir, exist_ok=True)
os.makedirs(graph_dir, exist_ok=True)

# 预定义时间点
required_dates = [
    "2022-07-01", "2022-08-01", "2022-09-01", 
    "2023-01-01", "2023-02-01", "2023-03-01", 
    "2023-04-01", "2023-05-01", "2023-06-01"
]

# 将日期转为 pandas 日期格式
required_dates = pd.to_datetime(required_dates)

# 用来记录 是否有date异常的category，false表示没有
flag = False

# 遍历每个 JSON 文件
for filename in os.listdir(input_dir):
    if filename.endswith(".json"):
        filepath = os.path.join(input_dir, filename)
        
        # 读取 JSON 数据
        with open(filepath, 'r', encoding='utf-8') as file:
            data = json.load(file)

        # 转换为 DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])

        # 检查时间范围
        if not df['date'].isin(required_dates).all():
            with open(os.path.join(log_dir, 'err.log'), 'a', encoding='utf-8') as log_file:
                log_file.write(f"Error in file {filename}: 这个category的日期不满足指定的格式： 2022的7-9 三个月，2023的1-6 六个月。\n")
            flag = True
            continue

        # 提取日期和库存量
        x = (df['date'] - pd.Timestamp("2022-07-01")).dt.days  # 转为天数
        y = df['inventory'].values

        # 三次样条插值
        cs = CubicSpline(x, y)

        # 计算从2022-07到2023-06的完整12个月的日期
        full_months = pd.date_range(start="2022-07-01", end="2023-06-01", freq='MS')
        x_full = (full_months - pd.Timestamp("2022-07-01")).days

        # 插值计算
        interpolated_values = cs(x_full)

        # 准备输出数据
        interpolated_data = [
            {"date": date.strftime('%Y-%m-%d'), "inventory": int(round(inv))}
            for date, inv in zip(full_months, interpolated_values)
        ]

        # 保存到指定目录
        output_filepath = os.path.join(output_dir, f"interpolation_{filename[18:-5]}.json")
        with open(output_filepath, 'w', encoding='utf-8') as output_file:
            json.dump(interpolated_data, output_file, ensure_ascii=False, indent=4)
        
        print(f"Processed and saved: {output_filepath}")

        # 使用更密集的点进行插值
        x_dense = np.linspace(x.min(), x.max(), 500)  # 增加插值点的数量
        y_dense = cs(x_dense)

        # 绘图
        plt.figure(figsize=(10, 6))

        # 原始数据点
        plt.scatter(x, y, color='blue', label='Original Data', zorder=5)

        # 插值数据点
        plt.scatter(x_full, interpolated_values, color='red', label='Interpolated Data', zorder=5)

        # 绘制样条曲线
        plt.plot(x_dense, y_dense, color='green', label='Cubic Spline Curve', zorder=3)

        # 设置图表标题和标签
        plt.title(f"Cubic Spline Interpolation for {filename[18:-5]}")
        plt.xlabel("Days since 2022-07-01")
        plt.ylabel("Inventory")
        plt.legend()
        plt.grid()

        # 保存图形
        graph_filepath = os.path.join(graph_dir, f"graph_category{filename[18:-5]}.png")
        plt.savefig(graph_filepath)
        plt.close()
        
        print(f"Graph saved: {graph_filepath}")


if (flag==False):
    with open(os.path.join(log_dir, 'err.log'), 'a', encoding='utf-8') as log_file:
        log_file.write(f"Congrats！ 所有的category的库存数据 都是严格遵守相同的日期格式， 即： 2022的7-9， 2023的1-6\n")
